package tool.workflow.fsSoft

import akka.stream.Materializer
import akka.stream.scaladsl.Sink
import akka.util.ByteString
import command.CommandExec
import implicits.Implicits._
import org.zeroturnaround.zip.ZipUtil
import play.api.libs.ws.WSClient
import shared.Pojo._
import tool.Pojo._
import shared.Pojo._
import tool.plot._
import tool.softT._
import tool.workflow.fsSoft.AllFsSoft.FSSoftBT
import tool.{FeatureSelectionTool, FormTool, Tool, plot}
import utils.Utils

import java.io.{File, FilenameFilter}
import scala.collection.immutable.SeqMap
import scala.concurrent.ExecutionContext.Implicits.global
import scala.concurrent.Future
import scala.concurrent.duration._
import scala.language.postfixOps
import dao._
import org.apache.commons.io.FileUtils
import play.api.i18n.Messages
import shared.pojo.WorkflowMission.EachDiagnoseAnaData
import shared.tool.DefaultPickle
import shared.{SharedTool, VarTool}
import tool.workflow.soft.AllSubSoft

import java.nio.file.Files

/** Created by yz on 3/12/2020
  */
object FSDiagnoseAna extends FSSoftBT with AllSoftBT with DiagnoseAnaSoftT {

  override def getFileInfos(workspaceDir: File, index: Int)(implicit
      matrixFileDataDao: MatrixFileDataDao,
      messages: Messages
  ): List[FileInfoData] = {
    val initInfos = {
      val kind = "singleMet"
      List(
        FileInfoData("Single_Met_ROC.csv", messages("singleMetROC"), kind = kind),
        FileInfoData("Single_Met_ROC.pdf", messages("singleMetROCCurve"), kind = kind)
      )
    }
    val data = FeatureSelectionTool.getFsData(workspaceDir).diagnoseAna(index)
    val rfMap = if (data.runRf) {
      SeqMap("RF" -> messages("rfModel"))
    } else SeqMap.empty
    val gbMap = if (data.runGb) {
      SeqMap("GB" -> messages("gbModel"))
    } else SeqMap.empty
    val lrMap = if (data.runLr) {
      SeqMap("LR" -> messages("lrModel"))
    } else SeqMap.empty
    val svmMap = if (data.runSvm) {
      SeqMap("SVM" -> messages("svmModel"))
    } else SeqMap.empty
    val enMap = if (data.runEn) {
      SeqMap("EN" -> messages("enModel"))
    } else SeqMap.empty
    val map = rfMap ++ gbMap ++ lrMap ++ svmMap ++ enMap
    val infos = map.flatMap { case (prefix, modelName) =>
      val kind = prefix.toLowerCase
      val targetDir = getDiagnoseWorkspaceDir(workspaceDir)
      val subDir = getSubDir(targetDir, prefix)
      val relPath = subDir.getRelPath(workspaceDir)
      val infos = List(
        FileInfoData(
          s"${prefix}_Prediction.csv",
          messages("predictionArg", s"${modelName}"),
          kind = kind,
          fileRelPath = relPath
        ),
        FileInfoData(
          s"${prefix}_Prediction_Summary.csv",
          messages("predictionSummaryArg", s"${modelName}"),
          kind = kind,
          fileRelPath = relPath
        ),
        FileInfoData(
          s"${prefix}_VarImp.csv",
          messages("varImpArg", s"${modelName}"),
          kind = kind,
          fileRelPath = relPath
        ),
        FileInfoData(s"ROC_Curve.pdf", messages("rOCCurveArg", s"${modelName}"), kind = kind, fileRelPath = relPath),
        FileInfoData(
          s"ROC_Curve_Data.csv",
          messages("rOCCurveDataArg", s"${modelName}"),
          kind = kind,
          fileRelPath = relPath
        ),
        FileInfoData(s"PR_Curve.pdf", messages("pRCurveArg", s"${modelName}"), kind = kind, fileRelPath = relPath),
        FileInfoData(
          s"PR_Curve_Data.csv",
          messages("pRCurveDataArg", s"${modelName}"),
          kind = kind,
          fileRelPath = relPath
        )
      )
      val cvInfos = if (data.hasCv && VarTool.cvPrefixs.contains(prefix)) {
        val fileRelPath = s"${relPath}/${targetDir.rfDir.cvDir.getName}"
        List(
          FileInfoData(
            s"${prefix}_Prediction_CV.csv",
            messages("predictionCV"),
            kind = kind,
            fileRelPath = fileRelPath
          ),
          FileInfoData(
            s"${prefix}_Prediction_Summary_CV_Overall.csv",
            messages("predictionSummaryCV"),
            kind = kind,
            fileRelPath = fileRelPath
          ),
          FileInfoData(s"ROC_Curve_CV_Overall.pdf", messages("rOCCurveCv"), kind = kind, fileRelPath = fileRelPath),
          FileInfoData(
            s"ROC_Curve_Data_CV_Overall.csv",
            messages("rOCCurveCvData"),
            kind = kind,
            fileRelPath = fileRelPath
          ),
          FileInfoData(
            s"ROC_Curve_CV_Overall+Each_Fold.pdf",
            messages("rocCurveCvEach"),
            kind = kind,
            fileRelPath = fileRelPath
          ),
          FileInfoData(
            s"ROC_Curve_Data_CV_Overall+Each_Fold.csv",
            messages("rocCurveDataCvEach"),
            kind = kind,
            fileRelPath = fileRelPath
          )
        )
      } else List.empty
      infos ::: cvInfos
    }.toList
    initInfos ::: infos
  }

  def produceSubConfigFile(workspaceDir: File, data: EachDiagnoseAnaData) = {
    val subConfigFile = new File(workspaceDir, "calculate_config.json")
    DefaultPickle.write(data).toFile(subConfigFile)
  }

  override def produceInput(workspaceDir: File, resultDir: File, index: Int, parentNodeStr: String, id: String)(implicit
      dao: MyDao,
      messages: Messages
  ): Unit = {
    val data = FeatureSelectionTool.getFsData(workspaceDir.getParentFile).diagnoseAna(index)

    val originalGroupFile = new File(workspaceDir.getParentFile, "group.txt")
    val groupFile = Tool.getGroupFile(workspaceDir)
    if (data.group.isBlank) {
      groupFile.fileCopyToDir(workspaceDir)
    } else {
      val groups = data.group.mySplit(" vs ")
      originalGroupFile.txtLines.lineSeqMap
        .filter { map =>
          groups.contains(map("ClassNote"))
        }
        .csvLines
        .toTxtFile(groupFile)
    }

    FeatureSelectionTool.produceSubInputFile(workspaceDir, parentNodeStr)
    produceSubConfigFile(workspaceDir, data)
    val rawDataFile = new File(workspaceDir, "AllMet_Raw.txt")
    val userIdR = "/user/(\\d+)/.*$".r
    val userId = userIdR
      .findFirstMatchIn(workspaceDir.unixPath)
      .map { matcher =>
        matcher.group(1)
      }
      .get
      .toInt
    if (data.hasExternal) {
      val testFile = new File(workspaceDir, "original_test.txt")
      Tool.produceRawMetaboliteFile(userId, data.externalFile, testFile)
      Tool.produceRawGroupFile(
        userId,
        data.externalGroupFile,
        testFile,
        new File(workspaceDir, "original_test_group.txt")
      )
    }
  }

  def getSubDir(diagnoseDir: DiagnoseDir, prefix: String) = {
    prefix match {
      case "RF" => diagnoseDir.rfDir.self
      case "GB" => diagnoseDir.gbDir.self
      case "LR" => diagnoseDir.lrDir.self
      case "SVM" => diagnoseDir.svmDir
      case "EN" => diagnoseDir.enDir
      case _ => diagnoseDir.self
    }
  }

  def getDiagnoseMap(data: EachDiagnoseAnaData) = {
    val rfMap = if (data.runRf) {
      SeqMap("RF" -> "fs_rf")
    } else SeqMap.empty
    val gbMap = if (data.runGb) {
      SeqMap("GB" -> "fs_gb")
    } else SeqMap.empty
    val lrMap = if (data.runLr) {
      SeqMap("LR" -> "fs_logistic")
    } else SeqMap.empty
    val svmMap = if (data.runSvm) {
      SeqMap("SVM" -> "fs_svm")
    } else SeqMap.empty
    val enMap = if (data.runEn) {
      SeqMap("EN" -> "fs_en")
    } else SeqMap.empty
    rfMap ++ gbMap ++ lrMap ++ svmMap ++ enMap
  }

  override def getCommandData(workspaceDir: File, resultDir: File, index: Int): List[CommandData] = {
    val fsData = FeatureSelectionTool.getFsData(workspaceDir.getParentFile)
    val data = fsData.diagnoseAna(index)
    val indexMap = SharedTool.getIndexMap(fsData.preprocess.datas.map(_.kind))
    val preprocessCommands = if (data.hasExternal) {
      if (data.externalPreprocess == "T") {
        val i = fsData.preprocess.datas.size
        val beforeCommand =
          s"""
             |cp original_test.txt 0.result.txt
             |cp group.txt original_group.txt
             |cp original_test_group.txt group.txt
             |""".stripMargin
        val commands = fsData.preprocess.datas.map(_.kind).zipWithIndex.map { kindI =>
          val (kind, i) = kindI
          val curSoft = AllSubSoft.kindSoftMap(kind)
          val index = indexMap(i)
          val si = List(i + 1, index).mkString(";")
          curSoft.getCommand(si)
        }
        val afterCommand =
          s"""
             |cp ${i}.result.txt pre_test.txt
             |cp original_group.txt group.txt
             |cp original_test_group.txt test_group.txt
             |""".stripMargin
        (beforeCommand +: commands) :+ afterCommand
      } else {
        val command =
          s"""
             |cp original_test_group.txt test_group.txt
             |cp original_test.txt pre_test.txt
             |""".stripMargin
        List(command)
      }
    } else {
      val command =
        s"""
           |cp ../preprocess_0/test.txt pre_test.txt
           |""".stripMargin
      List(command)
    }

    val solutionFile = new File(Tool.dataDir, "color_solution/color_solution_1.txt")
    val groupFile = new File(workspaceDir, "group.txt")
    val configJsonFile = Tool.getCalculateConfigJsonFile(workspaceDir)
    val dataFile = Tool.getAllMetRawFile(workspaceDir)
    val standardFile = Tool.getStandardFile(workspaceDir)
    val command =
      s"""
         |${"Rscript".condaPath()} ${Tool.rPath.unixPath}/standard.R --i ${new File(
          workspaceDir,
          "pre_test.txt"
        ).unixPath} --o  ${new File(workspaceDir, "test.txt").unixPath}
         |${"Rscript".condaPath()} ${new File(Tool.rPath, "roc_calculate.R").unixPath}
         |${"Rscript".condaPath()} ${new File(Tool.rPath, "roc_plot.R").unixPath}
           """.stripMargin
    val hasCv = data.hasCv
    val cvSplitCommand = if (hasCv) {
      s"""
         |${"Rscript".condaPath()} ${new File(Tool.rPath, "cv_split.R").unixPath} --i ${standardFile.unixPath}
         |""".stripMargin
    } else ""
    val scaleCommand =
      s"""
         |${"Rscript".condaPath()} ${Tool.rPath.unixPath}/standard.R
         |""".stripMargin
    val map = getDiagnoseMap(data)
    val otherCommandDatas = map.flatMap { case (prefix, rFileName) =>
      val diagnoseWorkspaceDir = getDiagnoseWorkspaceDir(workspaceDir)
      val subWorkspaceDir = getSubDir(diagnoseWorkspaceDir, prefix)
      val command =
        s"""
           |${"Rscript".condaPath()} ${new File(
            Tool.rPath,
            s"${rFileName}.R"
          ).unixPath} --config ${configJsonFile.unixPath} --i ${standardFile.unixPath} --g ${groupFile.unixPath}
           |""".stripMargin
      val plotCommands = List(
        plot.Roc.getCommand(subWorkspaceDir, "",convert = false),
        plot.Pr.getCommand(subWorkspaceDir, "",convert = false)
      )
      val cvDir = if (hasCv && VarTool.cvPrefixs.contains(prefix)) {
        new File(subWorkspaceDir, "cv")
      } else subWorkspaceDir
      val cvCommands = if (hasCv) {
        val calculateCommand = prefix match {
          case "RF" =>
            s"""
               |${"Rscript".condaPath()} ${new File(
                Tool.rPath,
                "rf_cv.R"
              ).unixPath} --i ${standardFile.unixPath} --g ${groupFile.unixPath}
               |""".stripMargin
          case "GB" =>
            s"""
               |${"Rscript".condaPath()} ${new File(
                Tool.rPath,
                "gb_cv.R"
              ).unixPath} --i ${standardFile.unixPath} --g ${groupFile.unixPath}
               |""".stripMargin
          case "LR" =>
            s"""
               |${"Rscript".condaPath()} ${new File(
                Tool.rPath,
                "lr_cv.R"
              ).unixPath} --i ${standardFile.unixPath} --g ${groupFile.unixPath}
               |""".stripMargin
          case _ => ""
        }
        val plotCommand = if (VarTool.cvPrefixs.contains(prefix)) {
          s"""
             |${"Rscript".condaPath()} ${new File(Tool.rPath, "cv_roc_plot.R").unixPath} --pc ${RocCv.configFileName}
             |convert  -density 300 ROC_Curve_CV_Overall.pdf ROC_Curve_CV_Overall.png
             |${"Rscript".condaPath()} ${new File(
              Tool.rPath,
              "cv_roc_all_plot.R"
            ).unixPath} --pc ${RocCvEach.configFileName}
             |convert  -density 300 ROC_Curve_CV_Overall+Each_Fold.pdf ROC_Curve_CV_Overall+Each_Fold.png
             |""".stripMargin
        } else ""
        List(calculateCommand, plotCommand)
      } else List.empty
      List(CommandData(subWorkspaceDir, command::plotCommands), CommandData(cvDir, cvCommands))
    }.toList
    val inputMetabolites = Tool.inputMetabolites(workspaceDir)
    val commandDatas = if (inputMetabolites.nonEmpty) {
      List(CommandData(workspaceDir, preprocessCommands :+ command :+ scaleCommand :+ cvSplitCommand)) :::
        otherCommandDatas
    } else List(CommandData(workspace = workspaceDir, command = ""))
    commandDatas
  }

  case class CvPDir(self: File, cvDir: File)

  case class DiagnoseDir(self: File, rfDir: CvPDir, gbDir: CvPDir, lrDir: CvPDir, svmDir: File, enDir: File)

  def getCvPDir(dir: File) = {
    val cvDir = new File(dir, "cv")
    CvPDir(self = dir, cvDir = cvDir)
  }

  def getDiagnoseWorkspaceDir(workspaceDir: File) = {
    val rfDir = new File(workspaceDir, "rf")
    val rfCvPDir = getCvPDir(rfDir)
    val gbDir = new File(workspaceDir, "gb")
    val gbCvPDir = getCvPDir(gbDir)
    val lrDir = new File(workspaceDir, "lr")
    val lrCvPDir = getCvPDir(lrDir)
    val svmDir = new File(workspaceDir, "svm")
    val enDir = new File(workspaceDir, "en")
    DiagnoseDir(
      self = workspaceDir,
      rfDir = rfCvPDir,
      gbDir = gbCvPDir,
      lrDir = lrCvPDir,
      svmDir = svmDir,
      enDir = enDir
    )
  }

  def getResultCvPDir(dir: File) = {
    val cvDir = new File(dir, "Cross-Validation-Results")
    CvPDir(self = dir, cvDir = cvDir)
  }

  def getDiagnoseResultDir(resultDir: File) = {
    val rfDir = new File(resultDir, "Random_Forest")
    val rfCvPDir = getResultCvPDir(rfDir)
    val gbDir = new File(resultDir, "Gradient_Boosting")
    val gbCvPDir = getResultCvPDir(gbDir)
    val lrDir = new File(resultDir, "Logistic_Regression")
    val lrCvPDir = getResultCvPDir(lrDir)
    val svmDir = new File(resultDir, "Support_Vector_Machine")
    val enDir = new File(resultDir, "Elastic_Net")
    DiagnoseDir(self = resultDir, rfDir = rfCvPDir, gbDir = gbCvPDir, lrDir = lrCvPDir, svmDir = svmDir, enDir = enDir)
  }

  def cvPDirCopy(sourceDir: CvPDir, targetDir: CvPDir)(implicit fileNameFilter: FilenameFilter) = {
    sourceDir.self.dirCopy(targetDir.self)
    sourceDir.cvDir.dirCopy(targetDir.cvDir)
  }

  def diagnoseCopy(sourceDir: DiagnoseDir, targetDir: DiagnoseDir)(implicit fileNameFilter: FilenameFilter) = {
    sourceDir.self.dirCopy(targetDir.self)
    cvPDirCopy(sourceDir.rfDir, targetDir.rfDir)
    cvPDirCopy(sourceDir.gbDir, targetDir.gbDir)
    cvPDirCopy(sourceDir.lrDir, targetDir.lrDir)
    sourceDir.svmDir.dirCopy(targetDir.svmDir)
    sourceDir.enDir.dirCopy(targetDir.enDir)
  }

  override def resultDeal(workspaceDir: File, resultDir: File, index: Int)(implicit
      ws: WSClient,
      materializer: Materializer,
      messages: Messages
  ): Unit = {
    val fsData = FeatureSelectionTool.getFsData(workspaceDir.getParentFile)
    val data = fsData.diagnoseAna(index)
    val map = getDiagnoseMap(data)
    map.toList.foreach { case (prefix, rFileName) =>
      val diagnoseWorkspaceDir = getDiagnoseWorkspaceDir(workspaceDir)
      val subWorkspaceDir = getSubDir(diagnoseWorkspaceDir, prefix)
      val needConvertFileNamePrefixs = List("ROC_Curve", "PR_Curve")
      val convertCommands = Tool.fileConvertCommandsIfExist(subWorkspaceDir, needConvertFileNamePrefixs)
      CommandExec().exec { b =>
        CommandData(workspace = subWorkspaceDir, commands = convertCommands, runFileName = "convert.sh")
      }
    }

    super.resultDeal(workspaceDir, resultDir, index)
    val subResultDir = new File(resultDir, s"${Tool.getTrueOutDir(s"${nameStr}_${index}")}")
    val diagnoseWorkspaceDir = getDiagnoseWorkspaceDir(workspaceDir)
    val diagnoseResultDir = getDiagnoseResultDir(subResultDir)
    implicit val filenameFilter = Tool.getFileNameFilter(List(".json"))
    diagnoseCopy(diagnoseWorkspaceDir, diagnoseResultDir)
  }

  override def producePlotConfigFile(workspaceDir: File, index: Int): Any = {
    val data = FeatureSelectionTool.getFsData(workspaceDir.getParentFile).diagnoseAna(index)
    val rfPrefix = if (data.runRf) {
      List("RF")
    } else {
      List.empty
    }
    val gbPrefix = if (data.runGb) {
      List("GB")
    } else {
      List.empty
    }
    val lrPrefix = if (data.runLr) {
      List("LR")
    } else {
      List.empty
    }
    val svmPrefix = if (data.runSvm) {
      List("SVM")
    } else {
      List.empty
    }
    val enPrefix = if (data.runEn) {
      List("EN")
    } else {
      List.empty
    }
    val prefixs = rfPrefix ::: gbPrefix ::: lrPrefix ::: svmPrefix ::: enPrefix
    val diagnoseWorkspaceDir = getDiagnoseWorkspaceDir(workspaceDir)
    prefixs.foreach { prefix =>
      val subWorkspaceDir = getSubDir(diagnoseWorkspaceDir, prefix)
      plot.Roc.initConfigFile(subWorkspaceDir, "")
      plot.Pr.initConfigFile(subWorkspaceDir, "")
      if (VarTool.cvPrefixs.contains(prefix) && data.hasCv) {
        val cvDir = new File(subWorkspaceDir, diagnoseWorkspaceDir.rfDir.cvDir.getName)
        Tool.producePlotConfigFile(AllPlotData(), RocCv, cvDir)
        Tool.producePlotConfigFile(AllPlotData(), RocCvEach, cvDir)
      }
    }
  }

}
